#!/usr/bin/env python
# predict by hsblock
import networkx as nx
import argparse
import sys

parser = argparse.ArgumentParser(
    description='link prediction by Hierarchical Stochastic Block Models',
    epilog='(c) 2008-2013, Yongjin Park and Joel S. Bader, {ypark28, joel.bader}@jhu.edu'
)

parser.add_argument( 'treepickle', help='tree model file' )
parser.add_argument( '-t', '--train', default=None, help='training network for degree corrected model' )
args = parser.parse_args()

T = nx.read_gpickle( args.treepickle )
test_pairs = map(lambda x: x.strip().split()[:2], sys.stdin)

## get appropriate degree correction
if args.train is not None:
    try:        
        elist = map(lambda x: x.strip().split('\t'), open( args.train ))
    except IOError:
        print >> sys.stderr, "cannot read training network: %s" % args.train
        sys.exit(1)

    G = nx.Graph()
    try:
        for e in elist :
            if len(e) < 2 : continue
            if e[0] == e[1] : continue
            if len(e) == 2 : 
                w = 1.
            elif len(e) > 2 : 
                w = float(e[2])
            G.add_edge( e[0], e[1], weight=w )
    except ValueError:
        print >> sys.stderr, "invalide edge weight in network %s " % args.train
        sys.exit(1)

    deg = G.degree( weight='weight' )
    degsum = sum(deg.values())
else:
    deg = {}
    degsum = 1.


def tree_visit(r):

    nd = T.node[r]
    p = nd.get('edges',0.) / max(0.001, nd.get('total',0.))

    if nd.get('isleaf',False) :
        bottom = set( nd.get('subset',[]) )

        # predict within
        for u,v in test_pairs :
            if u in bottom and v in bottom:
                d_u = deg.get(u,1.)
                d_v = deg.get(v,1.)
                print '%s\t%s\t%.4f'%(u, v, p * d_u * d_v / degsum)

        return bottom

    children = T.neighbors(r)

    bot_dict = {}
    for i,c in enumerate(children) :
        bot_dict[i] = tree_visit( c )

    # between pairs
    m = len(children)
    for i in xrange(m) :
        b_i = bot_dict[i]
        for j in xrange(i+1,m) :
            b_j = bot_dict[j]
            for u,v in test_pairs :
                if (u in b_i and v in b_j) or (u in b_j and v in b_i):
                    d_u = deg.get(u,1.)
                    d_v = deg.get(v,1.)
                    print '%s\t%s\t%.4f'%(u,v,p * d_u * d_v / degsum)

    bottom = reduce( lambda a,x : a.union( x ), bot_dict.values(), set([]) )
    return bottom

tree_visit(0)

